package pl.com.qapps.datastore.base;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;

import pl.com.qapps.datastore.ApiUtils;
import pl.com.qapps.datastore.EntityStore;
import pl.com.qapps.datastore.IdStore;
import pl.com.qapps.datastore.Store;
import pl.com.qapps.datastore.TransactionCore;
import pl.com.qapps.datastore.base.FutureHelper.CumulativeAggregateFuture;

import com.google.appengine.api.datastore.AsyncDatastoreService;
import com.google.appengine.api.datastore.DatastoreAttributes;
import com.google.appengine.api.datastore.DatastoreFailureException;
import com.google.appengine.api.datastore.DatastoreServiceConfig;
import com.google.appengine.api.datastore.Entity;
import com.google.appengine.api.datastore.EntityTranslator;
import com.google.appengine.api.datastore.Index;
import com.google.appengine.api.datastore.Index.IndexState;
import com.google.appengine.api.datastore.Key;
import com.google.appengine.api.datastore.KeyRange;
import com.google.appengine.api.datastore.Transaction;
import com.google.appengine.api.datastore.TransactionOptions;
import com.google.appengine.api.search.GetRequest;
import com.google.appengine.api.search.PutResponse;
import com.google.appengine.api.utils.FutureWrapper;
import com.google.appengine.labs.repackaged.com.google.common.base.Pair;
import com.google.appengine.labs.repackaged.com.google.common.collect.AbstractIterator;
import com.google.apphosting.api.DatastorePb.GetResponse;
import com.google.apphosting.api.DatastorePb.PutRequest;
import com.google.storage.onestore.v3.OnestoreEntity.EntityProto;

public class AsyncDatastoreServiceBase extends BaseDatastoreServiceImpl
		implements AsyncDatastoreService {

	/**
	 * {@link EntityGroupGrouper} that groups entities by their entity groups.
	 */
	static final EntityGroupGrouper<IndexedItem<Entity>> ENTITY_GROUPER = new EntityGroupGrouper<IndexedItem<Entity>>() {
		@Override
		Key extractEntityGroupKey(IndexedItem<Entity> item) {
			return getEntityGroupKey(item.item.getKey());
		}
	};

	/**
	 * {@link EntityGroupGrouper} that groups keys by their entity groups.
	 */
	static final EntityGroupGrouper<Key> KEY_GROUPER = new EntityGroupGrouper<Key>() {
		@Override
		Key extractEntityGroupKey(Key key) {
			return getEntityGroupKey(key);
		}
	};

	private final ExecutorService dispatcher;

	public AsyncDatastoreServiceBase(DatastoreServiceConfig cfg, Store store,
			TransactionStack stack, ExecutorService dispatcher) {
		super(cfg, store, stack, dispatcher);
		this.dispatcher = dispatcher;
	}

	/***
	 * An aggregate future that uses an iterator to match results to requested
	 * elements.
	 * 
	 * @param <K>
	 *            response type
	 * @param <I>
	 *            type being iterated over
	 * @param <V>
	 *            result type
	 */
	private abstract static class IteratingAggregateFuture<K, I, V> extends
			CumulativeAggregateFuture<K, Pair<Iterator<I>, V>, V> {
		public IteratingAggregateFuture(Iterable<Future<K>> futures) {
			super(futures);
		}

		protected abstract V aggregate(K intermediateResult,
				Iterator<I> iterator, V result);

		protected abstract Iterator<I> initIterator();

		protected abstract V initResult();

		@Override
		protected final Pair<Iterator<I>, V> aggregate(K intermediateResult,
				Pair<Iterator<I>, V> result) {
			return Pair.of(result.first,
					aggregate(intermediateResult, result.first, result.second));
		}

		@Override
		protected V finalizeResult(Pair<Iterator<I>, V> result) {
			return result.second;
		}

		@Override
		protected final Pair<Iterator<I>, V> initIntermediateResult() {
			return Pair.of(initIterator(), initResult());
		}
	}

	/**
	 * Models an item and its associated index in some ordered collection.
	 * 
	 * @param <T>
	 *            The type of the item.
	 */
	static class IndexedItem<T> implements Comparable<IndexedItem<T>> {
		final T item;
		final int index;

		IndexedItem(T item, int index) {
			this.item = item;
			this.index = index;
		}

		@Override
		public int compareTo(IndexedItem other) {
			return Integer.valueOf(index).compareTo(other.index);
		}
	}

	/**
	 * {@link Iterable implementation that converts a Iterable<IndexedItem<T>>
	 * to a Iterable<T>.
	 */
	private static class UnwrappingIterable<T> implements Iterable<T> {
		private final Iterable<IndexedItem<T>> innerIterable;

		private UnwrappingIterable(Iterable<IndexedItem<T>> innerIterable) {
			this.innerIterable = innerIterable;
		}

		@Override
		public Iterator<T> iterator() {
			return new AbstractIterator<T>() {
				Iterator<IndexedItem<T>> inner = innerIterable.iterator();

				@Override
				protected T computeNext() {
					if (inner.hasNext()) {
						return inner.next().item;
					}
					endOfData();
					return null;
				}
			};
		}
	}

	/**
	 * {@link CumulativeAggregateFuture} implementation that operates on Futures
	 * that return List<IndexedItem<Key>>. The final result produced by this
	 * Future returns the Keys sorted by their indexes.
	 */
	private static class SortingAggregateFuture
			extends
			CumulativeAggregateFuture<List<IndexedItem<Key>>, List<IndexedItem<Key>>, List<Key>> {
		private SortingAggregateFuture(
				Iterable<Future<List<IndexedItem<Key>>>> futures) {
			super(futures);
		}

		@Override
		protected List<IndexedItem<Key>> initIntermediateResult() {
			return new ArrayList<IndexedItem<Key>>();
		}

		@Override
		protected List<IndexedItem<Key>> aggregate(
				List<IndexedItem<Key>> intermediateResult,
				List<IndexedItem<Key>> result) {
			intermediateResult.addAll(result);
			return intermediateResult;
		}

		@Override
		protected List<Key> finalizeResult(
				List<IndexedItem<Key>> unorderedResult) {
			Collections.sort(unorderedResult);
			List<Key> orderedResult = new ArrayList<Key>(unorderedResult.size());
			for (IndexedItem<Key> key : unorderedResult) {
				orderedResult.add(key.item);
			}
			return orderedResult;
		}
	}

	@Override
	public Future<KeyRange> allocateIds(String kind, long num) {
		return allocateIds(null, kind, num);
	}

	@Override
	public Future<KeyRange> allocateIds(final Key parent, final String kind,
			final long num) {
		if (parent != null && !parent.isComplete())
			throw new IllegalArgumentException("parent cannot be incomplete");
		if (kind == null)
			throw new IllegalArgumentException("kind cannot be null");
		if (num < 1 || num > Long.MAX_VALUE - 1)
			throw new IllegalArgumentException("num must be between 0 and "
					+ (Long.MAX_VALUE - 1));

		Future<Long> idsFuture = dispatcher.submit(new Callable<Long>() {

			@Override
			public Long call() throws Exception {
				IdStore ids = store.getIdStore();
				return ids.next(parent, num);
			}

		});

		return new FutureWrapper<Long, KeyRange>(idsFuture) {

			@Override
			protected KeyRange wrap(Long lastId) throws Exception {
				return new KeyRange(parent, kind, lastId - num + 1, lastId);
			}

			@Override
			protected Throwable convertException(Throwable cause) {
				return cause;
			}

		};

		// AlocateIdsRequest req = new AlocateIdsRequest();
		// req.setAmount(num);
		//
		// return new FutureWrapper<Long, KeyRange>(runInTransaction(req)) {
		//
		// @Override
		// protected Throwable convertException(Throwable cause) {
		// return cause;
		// }
		//
		// @Override
		// protected KeyRange wrap(Long lastId) throws Exception {
		// return new KeyRange(parent, kind, lastId.longValue() - num + 1,
		// lastId);
		// }
		//
		// };
	}

	@Override
	public Future<Transaction> beginTransaction() {
		return beginTransaction(null);
	}

	@Override
	public Future<Transaction> beginTransaction(final TransactionOptions options) {

		return dispatcher.submit(new Callable<Transaction>() {

			@Override
			public Transaction call() throws Exception {
				TransactionCore core = store.getTransactionCore();
				return core.beginTransaction(
						txnProvider,
						dispatcher,
						options != null ? options : TransactionOptions.Builder
								.withDefaults());
			}

		});
	}

	@Override
	public Future<Void> delete(Key... keys) {
		return delete(null, Arrays.asList(keys));
	}

	@Override
	public Future<Void> delete(Iterable<Key> keys) {
		return delete(null, keys);
	}

	@Override
	public Future<Void> delete(Transaction txn, Key... keys) {
		return delete(txn, Arrays.asList(keys));
	}

	@Override
	public Future<Void> delete(final Transaction txn, final Iterable<Key> keys) {
		verifyKeys(keys);

		final EntityStore entities = store.getEntityStore();

		Future<Void> result = dispatcher.submit(new Callable<Void>() {

			@Override
			public Void call() throws Exception {
				for (Key key : keys) {
					entities.delete(key, txn);
				}
				return null;
			}

		});

		return new FutureWrapper<Void, Void>(result) {

			@Override
			protected Void wrap(Void nothing) throws Exception {
				return nothing;
			}

			@Override
			protected Throwable convertException(Throwable cause) {
				throw new DatastoreFailureException(cause.getMessage());
			}

		};

		// DeleteRequest req = new DeleteRequest();
		// req.setTxn(txn);
		// req.setKeys(keys);
		//
		// return runInTransaction(req);
	}

	@Override
	public Future<Entity> get(Key key) {
		return get(null, key);
	}

	@Override
	public Future<Map<Key, Entity>> get(Iterable<Key> keys) {
		return get(null, keys);
	}

	@Override
	public Future<Entity> get(Transaction txn, final Key key) {
		return new FutureWrapper<Map<Key, Entity>, Entity>(get(txn,
				Arrays.asList(key))) {

			@Override
			protected Throwable convertException(Throwable cause) {
				return cause;
			}

			@Override
			protected Entity wrap(Map<Key, Entity> map) throws Exception {
				return map.get(key);
			}
		};
	}

	@Override
	public Future<Map<Key, Entity>> get(final Transaction txn,
			final Iterable<Key> keys) {

		if (keys == null) {
			throw new NullPointerException("keys cannot be null");
		}

		if (txn == null
				&& datastoreServiceConfig.getMaxEntityGroupsPerRpc() != null
		// && datastoreServiceConfig.getReadPolicy().getConsistency() == STRONG
		// && getDatastoreType() == HIGH_REPLICATION
		) {
			Collection<List<Key>> keysByEntityGroup = KEY_GROUPER
					.getItemsByEntityGroup(keys);
			if (keysByEntityGroup.size() > 1) {
				return doBatchGetByEntityGroups(keysByEntityGroup);
			}
		}
		return doBatchGetBySize(txn, keys);

		// final int count = verifyKeys(keys);
		//
		// return dispatcher.submit(new Callable<Map<Key, Entity>>() {
		//
		// @Override
		// public Map<Key, Entity> call() throws Exception {
		// Map<Key, Entity> results = new HashMap<>(count);
		// EntityStore entities = store.getEntityStore();
		//
		// for (Key key : keys) {
		// Entity fetched = entities.get(key, txn);
		// results.put(key, fetched);
		// }
		//
		// return results;
		// }
		//
		// });
	}

	/**
	 * Executes a batch get, possibly by splitting into multiple rpcs to keep
	 * each rpc smaller than the maximum size.
	 * 
	 * @param txn
	 *            The transaction in which to execute the batch get. Can be
	 *            null.
	 * @param keys
	 *            The {@link Key keys} of the entities to fetch.
	 * 
	 * @return A {@link Future} that provides the results of the operation.
	 */
	private Future<Map<Key, Entity>> doBatchGetBySize(final Transaction txn,
			final Iterable<Key> keys) {

		ApiUtils.ensureTxnActive(txn);

		// if (datastoreServiceConfig.getReadPolicy().getConsistency() ==
		// EVENTUAL) {
		// baseReq.setFailoverMs(ARBITRARY_FAILOVER_READ_MS);
		// baseReq.setStrong(false);
		// }

		final List<Future<Iterable<Entity>>> futures = new ArrayList<>();
		for (Key key : keys) {
			if (!key.isComplete()) {
				throw new IllegalArgumentException(key + " is incomplete.");
			}
			// TODO split get queries by size
			// Reference ref = KeyTranslator.convertToPb(key);

			// int encodedKeySize =
			// ProtoConverter.stringSize(ref.encodingSize()) + 1;
			// if (datastoreServiceConfig.exceedsReadLimits(req.keySize() + 1,
			// encodedReqSize + encodedKeySize)) {
			// futures.add(makeAsyncCall(apiConfig, "Get", req,
			// new GetResponse()));
			// encodedReqSize = baseEncodedReqSize;
			// req = baseReq.clone();
			// }

			// encodedReqSize += encodedKeySize;
			// req.addKey(ref);
		}

		// if (req.keySize() > 0) {
		futures.add(dispatcher.submit(new Callable<Iterable<Entity>>() {

			@Override
			public Iterable<Entity> call() throws Exception {
				List<Entity> results = new ArrayList<>();
				EntityStore entities = store.getEntityStore();

				for (Key key : keys) {
					Entity fetched = entities.get(key, txn);
					results.add(fetched);
				}

				return results;
			}

		}));
		// original implementation
		// futures.add(makeAsyncCall(apiConfig, "Get", req, new
		// GetResponse()));
		// }

		return registerInTransaction(
				txn,
				new IteratingAggregateFuture<Iterable<Entity>, Key, Map<Key, Entity>>(
						futures) {
					@Override
					protected Map<Key, Entity> initResult() {
						return new HashMap<Key, Entity>();
					}

					@Override
					protected Iterator<Key> initIterator() {
						return keys.iterator();
					}

					@Override
					protected Map<Key, Entity> aggregate(
							Iterable<Entity> response,
							Iterator<Key> keyIterator, Map<Key, Entity> results) {
						for (Entity entity : response) {
							Key key = keyIterator.next();
							results.put(key, entity);
						}
						return results;
					}
				});
	}

	/**
	 * Executes a batch get by executing multiple rpcs in parallel.
	 * 
	 * @param keysByEntityGroup
	 *            A {@link Collection} of {@link List Lists} where all keys in
	 *            each list belong to the same entity group.
	 * 
	 * @return A {@link Future} that provides the results of all the get() rpcs.
	 */
	private Future<Map<Key, Entity>> doBatchGetByEntityGroups(
			Collection<List<Key>> keysByEntityGroup) {
		List<Future<Map<Key, Entity>>> subFutures = new ArrayList<Future<Map<Key, Entity>>>();
		List<Key> keysToGet = new ArrayList<Key>();
		int numEntityGroups = 0;
		for (List<Key> keysInGroup : keysByEntityGroup) {
			keysToGet.addAll(keysInGroup);
			numEntityGroups++;
			if (numEntityGroups == datastoreServiceConfig
					.getMaxEntityGroupsPerRpc()) {
				subFutures.add(doBatchGetBySize(null, keysToGet));
				keysToGet = new ArrayList<Key>();
				numEntityGroups = 0;
			}
		}
		if (!keysToGet.isEmpty()) {
			subFutures.add(doBatchGetBySize(null, keysToGet));
			// original implementation
			// futures.add(makeAsyncCall(apiConfig, "Put", req, new
			// PutResponse()));
		}
		return new CumulativeAggregateFuture<Map<Key, Entity>, Map<Key, Entity>, Map<Key, Entity>>(
				subFutures) {
			@Override
			protected Map<Key, Entity> initIntermediateResult() {
				return new HashMap<Key, Entity>();
			}

			@Override
			protected Map<Key, Entity> aggregate(
					Map<Key, Entity> intermediateResult, Map<Key, Entity> result) {
				intermediateResult.putAll(result);
				return intermediateResult;
			}

			@Override
			protected Map<Key, Entity> finalizeResult(Map<Key, Entity> result) {
				return result;
			}
		};
	}

	@Override
	public Future<DatastoreAttributes> getDatastoreAttributes() {
		throw new UnsupportedOperationException("not supported yet");
	}

	@Override
	public Future<Map<Index, IndexState>> getIndexes() {
		throw new UnsupportedOperationException("not supported yet or ever");
	}

	@Override
	public Future<Key> put(Entity entity) {
		GetOrCreateTransactionResult result = getOrCreateTransaction();
		return put(result.getTransaction(), entity);
	}

	@Override
	public Future<Key> put(Transaction txn, Entity entity) {
		return new FutureWrapper<List<Key>, Key>(
				put(txn, Arrays.asList(entity))) {
			@Override
			protected Key wrap(List<Key> keys) throws Exception {
				return keys.get(0);
			}

			@Override
			protected Throwable convertException(Throwable cause) {
				return cause;
			}
		};

	}

	@Override
	public Future<List<Key>> put(Iterable<Entity> entities) {
		GetOrCreateTransactionResult result = getOrCreateTransaction();
		return put(result.getTransaction(), entities);
	}

	/*
	 * (non-Javadoc)
	 * 
	 * @see
	 * com.google.appengine.api.datastore.AsyncDatastoreService#put(com.google
	 * .appengine.api.datastore.Transaction, java.lang.Iterable)
	 */
	@Override
	public Future<List<Key>> put(final Transaction txn,
			Iterable<Entity> entities) {

		if (txn == null
				&& datastoreServiceConfig.getMaxEntityGroupsPerRpc() != null) {
			List<IndexedItem<Entity>> indexedEntities = new ArrayList<IndexedItem<Entity>>();
			int index = 0;
			for (Entity entity : entities) {
				indexedEntities.add(new IndexedItem<Entity>(entity, index++));
			}
			Collection<List<IndexedItem<Entity>>> entitiesByEntityGroup = ENTITY_GROUPER
					.getItemsByEntityGroup(indexedEntities);
			if (entitiesByEntityGroup.size() > 1) {
				return doBatchPutByEntityGroups(entitiesByEntityGroup);
			}
		}
		return doBatchPutBySize(txn, entities);

		// TODO we don't implicitly allocate ids anymore
		// final List<Key> keys = new LinkedList<>();
		// for (Entity entity : entities) {
		// keys.add(entity.getKey());
		// if (!entity.getKey().isComplete()) {
		// // entities = generateIds(entities);
		// keys.clear();
		// Map<Key, Entity> generated = generateKeys(entities);
		// keys.addAll(generated.keySet());
		// entities = generated.values();
		// break;
		// }
		// }
		//
		// final Iterable<Entity> validEntities = entities;
		//
		// Future<Void> result = dispatcher.submit(new Callable<Void>() {
		//
		// @Override
		// public Void call() throws Exception {
		// EntityStore entityStore = store.getEntityStore();
		// for (Entity entity : validEntities) {
		// entityStore.put(entity, txn);
		// }
		// return null;
		// }
		// });
		// return new FutureWrapper<Void, List<Key>>(result) {
		//
		// @Override
		// protected List<Key> wrap(Void key) throws Exception {
		// return keys;
		// }
		//
		// @Override
		// protected Throwable convertException(Throwable cause) {
		// // TODO only temporary in development
		// return cause;
		// // return new DatastoreFailureException(cause.getMessage());
		// }
		// };
	}

	/**
	 * Executes a batch put, possibly by splitting into multiple calls to keep
	 * each call smaller than the maximum size.
	 * 
	 * @param txn
	 *            The transaction in which to execute the batch put. Can be
	 *            null.
	 * @param entities
	 *            The {@link Entity entities} to fetch.
	 * 
	 * @return A {@link Future} that provides the results of the operation.
	 */
	private Future<List<Key>> doBatchPutBySize(final Transaction txn,
			final Iterable<Entity> entities) {

		ApiUtils.ensureTxnActive(txn);

		PutRequest baseReq = new PutRequest();
		if (txn != null) {
			baseReq.setTransaction(localTxnToRemoteTxn(txn));
		}
		final int baseEncodedReqSize = baseReq.encodingSize();
		final List<Future<List<Key>>> futures = new ArrayList<Future<List<Key>>>();
		int encodedReqSize = baseEncodedReqSize;
		PutRequest req = baseReq.clone();
		for (Entity entity : entities) {
			EntityProto proto = EntityTranslator.convertToPb(entity);
			// int encodedEntitySize = Protocol.stringSize(proto.encodingSize())
			// + 1;
			// TODO spliting by size
			// if (datastoreServiceConfig.exceedsWriteLimits(req.entitySize() +
			// 1,
			// encodedReqSize + encodedEntitySize)) {
			// futures.add(makeAsyncCall(apiConfig, "Put", req,
			// new PutResponse()));
			// encodedReqSize = baseEncodedReqSize;
			// req = baseReq.clone();
			// }

			// encodedReqSize += encodedEntitySize;
			req.addEntity(proto);
		}

		if (req.entitySize() > 0) {
			futures.add(dispatcher.submit(new Callable<List<Key>>() {

				@Override
				public List<Key> call() throws Exception {
					List<Key> keys = new LinkedList<>();
					EntityStore entityStore = store.getEntityStore();
					for (Entity entity : entities) {
						entityStore.put(entity, txn);
						keys.add(entity.getKey());
					}
					return keys;
				}
			}));
		}

		return registerInTransaction(txn,
				new IteratingAggregateFuture<List<Key>, Entity, List<Key>>(
						futures) {
					@Override
					protected List<Key> initResult() {
						return new ArrayList<Key>();
					}

					@Override
					protected Iterator<Entity> initIterator() {
						return entities.iterator();
					}

					@Override
					protected List<Key> aggregate(List<Key> intermediateResult,
							Iterator<Entity> entitiesIterator,
							List<Key> keysInOrder) {
						// for (Reference reference : intermediateResult.keys())
						// {
						// Entity entity = entitiesIterator.next();
						// KeyTranslator.updateKey(reference, entity.getKey());
						// keysInOrder.add(entity.getKey());
						// }
						for (Key key : intermediateResult) {
							Entity entity = entitiesIterator.next();
							// TODO after put entity key needs to be updated
							// KeyTranslator.updateKey(reference,
							// entity.getKey());
							keysInOrder.add(entity.getKey());
						}
						return keysInOrder;
					}
				});
	}

	/**
	 * Executes a batch put by executing multiple rpcs in parallel.
	 * 
	 * @param entitiesByEntityGroup
	 *            A {@link Collection} of {@link List Lists} where all entities
	 *            in each list belong to the same entity group.
	 * 
	 * @return A {@link Future} that provides the results of all the put() rpcs.
	 */
	private Future<List<Key>> doBatchPutByEntityGroups(
			Collection<List<IndexedItem<Entity>>> entitiesByEntityGroup) {
		List<Future<List<IndexedItem<Key>>>> subFutures = new ArrayList<Future<List<IndexedItem<Key>>>>();
		List<IndexedItem<Entity>> entitiesToPut = new ArrayList<IndexedItem<Entity>>();
		int numEntityGroups = 0;
		for (List<IndexedItem<Entity>> indexedEntitiesInGroup : entitiesByEntityGroup) {
			entitiesToPut.addAll(indexedEntitiesInGroup);
			numEntityGroups++;
			if (numEntityGroups == datastoreServiceConfig
					.getMaxEntityGroupsPerRpc()) {
				assemblePutFuture(entitiesToPut, subFutures);
				numEntityGroups = 0;
			}
		}
		if (!entitiesToPut.isEmpty()) {
			assemblePutFuture(entitiesToPut, subFutures);
		}
		return new SortingAggregateFuture(subFutures);
	}

	/**
	 * Assembles a {@link Future} that puts the provided entities and then adds
	 * that Future to the provided {@link List}.
	 * 
	 * @param entitiesToPut
	 *            The entities to put.
	 * @param subFutures
	 *            The list of Futures.
	 */
	private void assemblePutFuture(List<IndexedItem<Entity>> entitiesToPut,
			List<Future<List<IndexedItem<Key>>>> subFutures) {
		final List<IndexedItem<Entity>> entitiesToPutCopy = new ArrayList<IndexedItem<Entity>>(
				entitiesToPut);
		Iterable<Entity> unwrappedEntitiesToPut = new UnwrappingIterable<Entity>(
				entitiesToPutCopy);
		Future<List<Key>> future = doBatchPutBySize(null,
				unwrappedEntitiesToPut);
		Future<List<IndexedItem<Key>>> indexedFuture = new FutureWrapper<List<Key>, List<IndexedItem<Key>>>(
				future) {
			@Override
			protected List<IndexedItem<Key>> wrap(List<Key> keys)
					throws Exception {
				List<IndexedItem<Key>> orderedKeys = new ArrayList<IndexedItem<Key>>(
						keys.size());
				int keyIndex = 0;
				for (Key key : keys) {
					orderedKeys.add(new IndexedItem<Key>(key, entitiesToPutCopy
							.get(keyIndex++).index));
				}
				return orderedKeys;
			}

			@Override
			protected Throwable convertException(Throwable cause) {
				return cause;
			}
		};
		subFutures.add(indexedFuture);
		entitiesToPut.clear();
	}

	@Deprecated
	private List<Entity> generateIds(Iterable<Entity> entities) {
		List<Entity> result = new LinkedList<>();
		for (Entity entity : entities) {
			if (!entity.getKey().isComplete()) {
				Key generated = FutureHelper.quietGet(
						allocateIds(entity.getParent(), entity.getKind(), 1L))
						.getStart();
				Entity newEntity = new Entity(generated);
				newEntity.setPropertiesFrom(entity);
				result.add(newEntity);
			} else
				result.add(entity);
		}
		return result;
		// Map<Pair<Key, String>, List<Entity>> toGenerate = new HashMap<>();
		// List<Entity> result = new LinkedList<>();
		// for (Entity entity : entities) {
		// if (!entity.getKey().isComplete()) {
		// Pair<Key, String> pair = Pair.of(entity.getParent(),
		// entity.getKind());
		// List<Entity> common = toGenerate.get(pair);
		// if (common == null) {
		// common = new LinkedList<>();
		// toGenerate.put(pair, common);
		// }
		// common.add(entity);
		// } else result.add(entity);
		// }
		// for(Entry<Pair<Key, String>, List<Entity>> entry :
		// toGenerate.entrySet()) {
		// for(Key key : allocateIds(entry.getKey().first,
		// entry.getKey().second,entry.getValue().size())) {
		//
		// }
		// }
	}

	// TODO i don't think that allocating id should be automatic, user should
	// explicitly do that before calling put
	private Map<Key, Entity> generateKeys(Iterable<Entity> entities) {
		Map<Key, Entity> result = new HashMap<>();
		for (Entity entity : entities) {
			if (!entity.getKey().isComplete()) {
				Key generated = FutureHelper.quietGet(
						allocateIds(entity.getParent(), entity.getKind(), 1L))
						.getStart();
				Entity newEntity = new Entity(generated);
				newEntity.setPropertiesFrom(entity);
				result.put(generated, newEntity);
			} else
				result.put(entity.getKey(), entity);
		}
		return result;
	}

	/**
	 * Register the provided future with the provided txn so that we know to
	 * perform a {@link java.util.concurrent.Future#get()} before the txn is
	 * committed.
	 * 
	 * @param txn
	 *            The txn with which the future must be associated.
	 * @param future
	 *            The future to associate with the txn.
	 * @param <T>
	 *            The type of the Future
	 * @return The same future that was passed in, for caller convenience.
	 */
	private <T> Future<T> registerInTransaction(Transaction txn,
			final Future<T> future) {
		if (txn != null) {
			txnProvider.addFuture(txn, future);
			return new FutureHelper.TxnAwareFuture<T>(future, txn, txnProvider);
		}
		return future;
	}

	private static int verifyKeys(Iterable<Key> keys) {
		int count = 0;
		for (Key key : keys) {
			if (key == null)
				throw new IllegalArgumentException("one of given keys is null");
			if (!key.isComplete()) {
				throw new IllegalArgumentException("key '" + key
						+ "' is incomplete");
			}
			count++;
		}
		return count;
	}

	/**
	 * A class that knows how to group items by entity group.
	 * 
	 * @param <T>
	 *            The type of item.
	 */
	abstract static class EntityGroupGrouper<T> {

		/**
		 * Arranges the given items by entity group.
		 * 
		 * @param items
		 *            The items to arrange by entity group.
		 * @return A {@link Collection} of {@link List Lists} where each
		 *         {@link List} contains all items belonging to the same entity
		 *         group.
		 */
		public Collection<List<T>> getItemsByEntityGroup(Iterable<T> items) {
			Map<Key, List<T>> entitiesByEntityGroup = new LinkedHashMap<Key, List<T>>();
			for (T item : items) {
				Key entityGroupKey = extractEntityGroupKey(item);
				List<T> entitiesInGroup = entitiesByEntityGroup
						.get(entityGroupKey);
				if (entitiesInGroup == null) {
					entitiesInGroup = new ArrayList<T>();
					entitiesByEntityGroup.put(entityGroupKey, entitiesInGroup);
				}
				entitiesInGroup.add(item);
			}
			return entitiesByEntityGroup.values();
		}

		static Key getEntityGroupKey(Key key) {
			Key curKey = key;
			while (curKey.getParent() != null) {
				curKey = curKey.getParent();
			}
			return curKey;
		}

		/**
		 * Given an item, extract the entity group key.
		 */
		abstract Key extractEntityGroupKey(T item);
	}

}
